﻿// --------------------------------------------------------------------------------------------------------------------
// <copyright file="PartitionedCloudQueue.cs" company="Microsoft">
// Copyright (c) Microsoft Corporation.  All rights reserved. 
// </copyright>
// --------------------------------------------------------------------------------------------------------------------

using System;
using System.Collections.Generic;
using System.Linq;
using System.Reflection;
using System.Threading;
using Microsoft.WindowsAzure.StorageClient;

namespace Microsoft.WindowsAzure.StorageClientExtensions
{   
    public class PartitionedCloudQueue
    {
        /// <summary>
        /// Wrapper of CloudQueueMessage which also keeps track of the parent Queue
        /// </summary>
        private class CloudQueueMessageWrapper : CloudQueueMessage
        {
            public CloudQueue OriginQueue;

            public CloudQueueMessageWrapper(byte[] content)
                : base(content)
            { }
           
            public static CloudQueueMessageWrapper FromCloudQueueMessage(CloudQueueMessage msg, CloudQueue originQ)
            {
                if (msg == null)
                {
                    return null;
                }

                CloudQueueMessageWrapper result = new CloudQueueMessageWrapper(msg.AsBytes);
                FieldInfo[] fields_of_class = msg.GetType().GetFields(BindingFlags.NonPublic | BindingFlags.Public | BindingFlags.Instance);

                foreach (FieldInfo fi in fields_of_class)
                {
                    fi.SetValue(result, fi.GetValue(msg));
                }


                result.OriginQueue = originQ;

                return result;
            }
        }
        
        private readonly List<CloudQueue> _partitions;
        private int _currentQueueIndex;
        private readonly CloudQueueClient _queueClient;
        private readonly String _queueBaseName;

        /// <summary>
        /// Method to randomly change the order of the elements in a sequence
        /// </summary>
        private static IEnumerable<T> ShuffleEnumerable<T>(IEnumerable<T> input)
        {
            Random rnd = new Random();
            return input.OrderBy(x => rnd.Next());
        }

        private int GetNextQueueIndex()
        {
            int signedres = Interlocked.Increment(ref _currentQueueIndex);
            uint unsignedres = (uint)signedres;
            return (int)(unsignedres % _partitions.Count);
        }

        /// <summary>
        /// Constructor of the PartitionedCloudQueue which allows to define number of partitions
        /// </summary>        
        public PartitionedCloudQueue(CloudQueueClient queueClient, string queueBaseName, int numberOfPartitions)            
        {
            _queueClient = queueClient;
            _queueBaseName = queueBaseName;
                        
            List<CloudQueue> partitions = new List<CloudQueue>();
            for (int i = 0; i < numberOfPartitions; i++)
            {
                CloudQueue q = queueClient.GetQueueReference(_queueBaseName + i);
                partitions.Add(q);
            }

            _partitions = new List<CloudQueue>(ShuffleEnumerable(partitions));
        }

        /// <summary>
        /// constructor of the PartitionedCloudQueue which allows to discover the already existing partitions
        /// There is assumption that queue already exists in Azure Storage
        /// </summary>        
        public PartitionedCloudQueue(CloudQueueClient queueClient, string queueBaseName)         
        {
            _queueClient = queueClient;
            _queueBaseName = queueBaseName;

            // discover how many partitions are there in the storage account
            List<CloudQueue> partitions = new List<CloudQueue>();
            partitions.Add(queueClient.GetQueueReference(_queueBaseName + 0));

            for (int index = 1; ; index++)
            {
                CloudQueue q = queueClient.GetQueueReference(_queueBaseName + index);
                if (q.Exists())
                {
                    partitions.Add(q);
                }
                else
                {
                    break;
                }
            }           

            _partitions = new List<CloudQueue>(ShuffleEnumerable(partitions));
        }
       
        #region CloudQueue implementation

        public void AddMessage(CloudQueueMessage message, TimeSpan timeToLive)
        {
            _partitions[GetNextQueueIndex()].AddMessage(message, timeToLive);
        }

        public void AddMessage(CloudQueueMessage message)
        {
            _partitions[GetNextQueueIndex()].AddMessage(message);
        }
               
        public IAsyncResult BeginAddMessage(CloudQueueMessage message, AsyncCallback callback, object state)
        {
            // select one of the underlying queues to work with
            CloudQueue queue = _partitions[GetNextQueueIndex()];
            NestedAsyncHelper wrapper = NestedAsyncHelper.WrapBeginParameters(callback, state, queue);
            return wrapper.WrapAsyncResult(queue.BeginAddMessage(message, NestedAsyncHelper.Callback, wrapper));
        }

        public void EndAddMessage(IAsyncResult asyncResult)
        {            
            CloudQueue queue = (CloudQueue)NestedAsyncHelper.GetExtraState(asyncResult);            
            queue.EndAddMessage(NestedAsyncHelper.UnwrapAsyncResult(asyncResult));
        }
      
        public CloudQueueMessage GetMessage()
        {
            CloudQueue originQ = _partitions[GetNextQueueIndex()];
            CloudQueueMessage msg = originQ.GetMessage();
            return CloudQueueMessageWrapper.FromCloudQueueMessage(msg, originQ);
        }

        public IEnumerable<CloudQueueMessage> GetMessages(int messageCount, TimeSpan visibilityTimeout)
        {
            CloudQueue originQ = _partitions[GetNextQueueIndex()];
            IEnumerable<CloudQueueMessage> msgs = originQ.GetMessages(messageCount, visibilityTimeout);
            return msgs.Select(msg => CloudQueueMessageWrapper.FromCloudQueueMessage(msg, originQ));
        }

        public IAsyncResult BeginGetMessages(int messageCount, TimeSpan visibilityTimeout, AsyncCallback callback, object state)
        {            
            CloudQueue queue = _partitions[GetNextQueueIndex()];
            NestedAsyncHelper wrapper = NestedAsyncHelper.WrapBeginParameters(callback, state, queue);
            return wrapper.WrapAsyncResult(queue.BeginGetMessages(messageCount, visibilityTimeout, NestedAsyncHelper.Callback, wrapper));
        }

        public IEnumerable<CloudQueueMessage> EndGetMessages(IAsyncResult asyncResult)
        {
            CloudQueue queue = (CloudQueue)NestedAsyncHelper.GetExtraState(asyncResult);
            IEnumerable<CloudQueueMessage> msgs = queue.EndGetMessages(NestedAsyncHelper.UnwrapAsyncResult(asyncResult));
            return msgs.Select(msg => CloudQueueMessageWrapper.FromCloudQueueMessage(msg, queue));            
        }

        public void DeleteMessage(CloudQueueMessage message)
        {
            if (message == null)
            {
                throw new ArgumentNullException("message");
            }

            CloudQueueMessageWrapper msg = (CloudQueueMessageWrapper)message;
            msg.OriginQueue.DeleteMessage(message);
        }

        public IAsyncResult BeginDeleteMessage(CloudQueueMessage message, AsyncCallback callback, object state)
        {
            if (message == null)
            {
                throw new ArgumentNullException("message");
            }

            CloudQueueMessageWrapper msg = (CloudQueueMessageWrapper)message;
            NestedAsyncHelper wrapper = NestedAsyncHelper.WrapBeginParameters(callback, state, msg.OriginQueue);

            return wrapper.WrapAsyncResult(msg.OriginQueue.BeginDeleteMessage(msg, NestedAsyncHelper.Callback, wrapper));          
        }

        public void EndDeleteMessage(IAsyncResult asyncResult)
        {
            CloudQueue queue = (CloudQueue)NestedAsyncHelper.GetExtraState(asyncResult);
            queue.EndDeleteMessage(NestedAsyncHelper.UnwrapAsyncResult(asyncResult));            
        }

        public int RetrieveApproximateMessageCount()
        {
            return _partitions.Select(o => o.RetrieveApproximateMessageCount()).Sum();
        }

        public bool CreateIfNotExist()
        {
            bool res=false;
            foreach (CloudQueue partition in _partitions)
            {
                res |= partition.CreateIfNotExist();
            }
            return res;
        }

        public bool Exists()
        {            
            return _partitions.All(o => o.Exists());
        }

        public void Delete()
        {
            foreach (CloudQueue partition in _partitions)
            {
                partition.Delete();
            }
        }

        #endregion
    }
}
